#!/usr/bin/env python
"""
将三元组数据(triples)从文本文件处理并添加实体类型标签，然后导入到Neo4j图数据库
文件格式: 实体A \t 关系 \t 实体B
例如: Cerberus    isA    banking_malware
"""

import logging
import os
import sys
from typing import Dict, List, Tuple, Set
from collections import defaultdict
import time
import random
from neo4j import GraphDatabase

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

class TripleProcessor:
    """处理三元组数据并添加实体类型标签的工具类"""
    
    # 定义实体类型
    ENTITY_TYPES = [
        'Malware',          # 恶意软件
        'Vulnerability',    # 漏洞
        'ThreatActor',      # 威胁行动者
        'Asset',            # 资产
        'Attack',           # 攻击
        'Tool',             # 工具
        'Campaign',         # 攻击活动
        'Technique',        # 攻击技术
        'Protocol',         # 协议
        'Software'          # 软件
    ]
    
    # 类型颜色映射（用于Neo4j显示）
    ENTITY_COLOR_MAP = {
        'Malware': '#FF5555',        # 红色
        'Vulnerability': '#FFA500',  # 橙色
        'ThreatActor': '#8B0000',    # 深红色
        'Asset': '#4682B4',          # 钢蓝色
        'Attack': '#9932CC',         # 紫色
        'Tool': '#008000',           # 绿色
        'Campaign': '#D2691E',       # 巧克力色
        'Technique': '#800080',      # 紫色
        'Protocol': '#20B2AA',       # 浅海绿
        'Software': '#6495ED'        # 矢车菊蓝
    }
    
    def __init__(self, batch_size=500, uri=None, user=None, password=None):
        """初始化处理器"""
        self.batch_size = batch_size
        self.neo4j_connected = False
        
        # 连接Neo4j数据库
        if uri and user and password:
            try:
                self.driver = GraphDatabase.driver(uri, auth=(user, password))
                self.driver.verify_connectivity()
                self.neo4j_connected = True
                logging.info(f"成功连接到Neo4j数据库: {uri}")
            except Exception as e:
                self.driver = None
                logging.error(f"Neo4j连接失败: {str(e)}")
                raise ConnectionError(f"无法连接到Neo4j数据库: {str(e)}")
        else:
            logging.error("未提供Neo4j连接信息")
            raise ValueError("必须提供Neo4j连接信息 (uri, user, password)")
        
        logging.info("三元组处理器初始化完成")
    
    def close(self):
        """关闭数据库连接"""
        if self.neo4j_connected and self.driver:
            self.driver.close()
            logging.info("Neo4j连接已关闭")
    
    def clear_database(self):
        """清空Neo4j数据库中的所有数据"""
        if not self.neo4j_connected:
            logging.error("未连接到Neo4j数据库，无法清空数据")
            return False
        
        try:
            with self.driver.session() as session:
                # 删除所有节点和关系
                session.run("MATCH (n) DETACH DELETE n")
                logging.info("成功清空Neo4j数据库中的所有数据")
            return True
        except Exception as e:
            logging.error(f"清空数据库失败: {str(e)}")
            return False
    
    def create_indexes(self):
        """创建必要的索引以加速导入和查询"""
        if not self.neo4j_connected:
            logging.error("未连接到Neo4j数据库，无法创建索引")
            return False
        
        try:
            with self.driver.session() as session:
                # 为不同类型的实体创建索引
                for entity_type in self.ENTITY_TYPES:
                    session.run(f"CREATE INDEX IF NOT EXISTS FOR (n:{entity_type}) ON (n.name)")
                
                # 为Entity创建索引
                session.run("CREATE INDEX IF NOT EXISTS FOR (n:Entity) ON (n.name)")
                
                logging.info("成功创建实体类型索引")
            return True
        except Exception as e:
            logging.error(f"创建索引失败: {str(e)}")
            return False
    
    def read_triples(self, file_path: str) -> List[Tuple[str, str, str]]:
        """从文件中读取三元组数据"""
        triples = []
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    try:
                        line = line.strip()
                        if not line or line.startswith('#'):
                            continue
                        
                        parts = line.split('\t')
                        if len(parts) != 3:
                            logging.warning(f"第{line_num}行格式不正确，跳过: {line}")
                            continue
                        
                        subject, predicate, obj = parts
                        triples.append((subject.strip(), predicate.strip(), obj.strip()))
                    except Exception as e:
                        logging.warning(f"处理第{line_num}行时出错: {str(e)}")
                        continue
        except Exception as e:
            logging.error(f"读取文件失败: {str(e)}")
        
        return triples
    
    def determine_entity_type(self, entity_name: str, relations: Dict[str, List[str]]) -> str:
        """根据实体名称和关系推断实体类型"""
        # 实体类型判断逻辑
        entity_name_lower = entity_name.lower()
        
        # 判断关系和实体名称特征来分配类型
        # 1. Malware类型判断
        if any(keyword in entity_name_lower for keyword in ["malware", "virus", "trojan", "ransomware", "spyware", "worm", "backdoor"]) or \
           any("isA" in rel and any(keyword in target.lower() for keyword in ["malware", "virus", "trojan"]) for rel, targets in relations.items() for target in targets):
            return "Malware"
        
        # 2. Vulnerability类型判断
        if any(keyword in entity_name_lower for keyword in ["vulnerability", "cve", "exploit", "weakness", "flaw", "bug"]) or \
           any("exploits" in rel for rel in relations.keys()):
            return "Vulnerability"
            
        # 3. ThreatActor类型判断
        if any(keyword in entity_name_lower for keyword in ["actor", "group", "apt", "hacker", "attacker", "threat"]) or \
           any("uses" in rel or "conducts" in rel or "sponsors" in rel for rel in relations.keys()):
            return "ThreatActor"
            
        # 4. Asset类型判断
        if any(keyword in entity_name_lower for keyword in ["server", "system", "network", "device", "host", "computer", "router", "firewall", "database"]):
            return "Asset"
            
        # 5. Attack类型判断
        if any(keyword in entity_name_lower for keyword in ["attack", "breach", "incident", "compromise"]):
            return "Attack"
            
        # 6. Tool类型判断
        if any(keyword in entity_name_lower for keyword in ["tool", "toolkit", "utility", "framework", "platform", "software"]) and not "malware" in entity_name_lower:
            return "Tool"
            
        # 7. Campaign类型判断
        if any(keyword in entity_name_lower for keyword in ["campaign", "operation", "mission", "activity"]):
            return "Campaign"
            
        # 8. Technique类型判断
        if any(keyword in entity_name_lower for keyword in ["technique", "tactic", "method", "procedure", "ttps", "approach"]):
            return "Technique"
            
        # 9. Protocol类型判断
        if any(keyword in entity_name_lower for keyword in ["protocol", "http", "https", "ftp", "ssh", "tcp", "udp", "ip", "dns"]):
            return "Protocol"
            
        # 10. Software类型判断
        if any(keyword in entity_name_lower for keyword in ["software", "program", "application", "app", "service"]) and not any(keyword in entity_name_lower for keyword in ["malware", "virus", "trojan"]):
            return "Software"
        
        # 如果通过名称无法确定，尝试通过关系推断
        if "isUsedBy" in relations or "targets" in relations:
            return "Tool"
            
        if "employsTechnique" in relations or "hasVulnerability" in relations:
            return "Campaign"
            
        # 如果无法确定，随机分配一个类型，但更倾向于常见类型
        weights = [0.3, 0.2, 0.15, 0.15, 0.1, 0.03, 0.02, 0.02, 0.02, 0.01]  # 权重按ENTITY_TYPES的顺序
        return random.choices(self.ENTITY_TYPES, weights=weights, k=1)[0]
    
    def process_triples(self, file_path: str) -> Dict[str, any]:
        """处理三元组数据并添加实体类型"""
        try:
            # 读取三元组
            triples = self.read_triples(file_path)
            logging.info(f"从文件中读取了 {len(triples)} 个三元组")
            
            # 收集所有唯一的实体和关系
            entities = set()
            relationships = set()
            
            # 构建实体关系图
            outgoing_relations = defaultdict(lambda: defaultdict(list))
            incoming_relations = defaultdict(lambda: defaultdict(list))
            
            for subject, predicate, obj in triples:
                entities.add(subject)
                entities.add(obj)
                relationships.add(predicate)
                
                # 记录实体的出向和入向关系
                outgoing_relations[subject][predicate].append(obj)
                incoming_relations[obj][predicate].append(subject)
            
            # 为每个实体确定一个类型
            entity_types = {}
            for entity in entities:
                entity_types[entity] = self.determine_entity_type(
                    entity, 
                    {**outgoing_relations.get(entity, {}), **incoming_relations.get(entity, {})}
                )
            
            # 构建节点和边的数据结构
            nodes = []
            links = []
            
            # 处理节点
            for entity in entities:
                nodes.append({
                    "id": entity,
                    "name": entity,
                    "type": entity_types[entity],
                    "outDegree": sum(len(targets) for targets in outgoing_relations[entity].values()),
                    "inDegree": sum(len(sources) for sources in incoming_relations[entity].values())
                })
            
            # 处理关系
            for subject, predicate, obj in triples:
                links.append({
                    "source": subject,
                    "target": obj,
                    "relation": predicate
                })
            
            # 清空Neo4j数据库
            if not self.clear_database():
                return {"success": False, "error": "清空Neo4j数据库失败"}
            
            # 创建索引
            self.create_indexes()
            
            # 导入数据到Neo4j
            if not self.import_to_neo4j(nodes, links):
                return {"success": False, "error": "导入数据到Neo4j失败"}
            
            # 生成类型统计信息
            type_stats = {}
            for entity_type in self.ENTITY_TYPES:
                type_stats[entity_type] = len([e for e in entity_types.values() if e == entity_type])
            
            # 处理统计信息
            stats = {
                "success": True,
                "total_triples": len(triples),
                "unique_entities": len(entities),
                "unique_relationships": len(relationships),
                "entity_type_stats": type_stats
            }
            
            logging.info(f"处理完成! 共有 {len(triples)} 个三元组，{len(entities)} 个唯一实体，{len(relationships)} 种关系类型")
            logging.info(f"实体类型统计: {type_stats}")
            logging.info("数据已成功导入到Neo4j数据库")
            
            return stats
        except Exception as e:
            logging.error(f"处理失败: {str(e)}")
            import traceback
            logging.error(traceback.format_exc())
            return {"success": False, "error": str(e)}
    
    def import_to_neo4j(self, nodes, links):
        """将处理好的数据导入到Neo4j数据库"""
        if not self.neo4j_connected:
            logging.error("未连接到Neo4j数据库，无法导入数据")
            return False
        
        try:
            # 1. 批量创建节点
            logging.info(f"开始批量创建 {len(nodes)} 个实体节点...")
            
            with self.driver.session() as session:
                # 首先配置不同标签的颜色
                for entity_type, color in self.ENTITY_COLOR_MAP.items():
                    # 设置Neo4j浏览器的可视化配置
                    config_query = f"""
                    CALL apoc.meta.graphSample()
                    YIELD relationships
                    WITH collect(relationships) AS rels
                    CALL db.labels() YIELD label
                    WITH label, rels
                    WHERE label = '{entity_type}'
                    CALL apoc.browser.setConfig('browser.visualization.nodeColorMapping', 
                       apoc.convert.toJson({{ '{entity_type}': '{color}' }})) YIELD value
                    RETURN value
                    """
                    try:
                        session.run(config_query)
                        logging.info(f"为 {entity_type} 节点设置颜色: {color}")
                    except Exception as e:
                        # 可能是APOC没有安装，或者Neo4j版本不支持
                        logging.warning(f"设置 {entity_type} 节点颜色失败，节点会使用默认颜色: {str(e)}")
                
                # 分批处理节点
                for i in range(0, len(nodes), self.batch_size):
                    batch = nodes[i:i+self.batch_size]
                    
                    # 为节点添加颜色属性
                    for node in batch:
                        node_type = node["type"]
                        if node_type in self.ENTITY_COLOR_MAP:
                            node["color"] = self.ENTITY_COLOR_MAP[node_type]
                    
                    # 为不同类型的实体分别创建节点
                    for entity_type in self.ENTITY_TYPES:
                        type_batch = [n for n in batch if n["type"] == entity_type]
                        if type_batch:
                            # 创建带有颜色属性的实体节点
                            type_query = f"""
                            UNWIND $batch AS node
                            MERGE (n:Entity:{entity_type} {{name: node.name}})
                            SET n.color = node.color,
                                n.outDegree = node.outDegree,
                                n.inDegree = node.inDegree,
                                n.type = node.type,
                                n.size = 10 + (node.outDegree + node.inDegree) / 2,  // 根据度数设置节点大小
                                n.created_at = datetime()
                            """
                            session.run(type_query, batch=type_batch)
                    
                    current_batch = (i // self.batch_size) + 1
                    total_batches = (len(nodes) + self.batch_size - 1) // self.batch_size
                    logging.info(f"创建实体节点: {current_batch}/{total_batches} 批完成")
            
            # 2. 批量创建关系
            logging.info(f"开始批量创建 {len(links)} 个关系...")
            
            # 按关系类型分组
            links_by_relation = defaultdict(list)
            for link in links:
                links_by_relation[link["relation"]].append((link["source"], link["target"]))
            
            with self.driver.session() as session:
                for relation, rel_pairs in links_by_relation.items():
                    # 格式化关系类型为Neo4j兼容的格式
                    relation_type = self._format_relation_name(relation)
                    total_batches = (len(rel_pairs) + self.batch_size - 1) // self.batch_size
                    
                    for i in range(0, len(rel_pairs), self.batch_size):
                        batch = rel_pairs[i:i+self.batch_size]
                        batch_data = [{"source": s, "target": t, "original": relation} for s, t in batch]
                        
                        # 创建关系，添加更多属性
                        query = f"""
                        UNWIND $batch AS pair
                        MATCH (s:Entity {{name: pair.source}}), (t:Entity {{name: pair.target}})
                        MERGE (s)-[r:{relation_type}]->(t)
                        ON CREATE SET 
                            r.created_at = datetime(), 
                            r.original_relation = pair.original,
                            r.weight = 1 + toInteger(rand() * 3),  // 随机权重，用于线条粗细
                            r.timestamp = timestamp()
                        """
                        session.run(query, batch=batch_data)
                        
                        current_batch = (i // self.batch_size) + 1
                        logging.info(f"创建 {relation_type} 关系: {current_batch}/{total_batches} 批完成")
            
            # 3. 设置Neo4j Browser可视化配置
            try:
                # 节点样式配置
                node_styles = {}
                for entity_type, color in self.ENTITY_COLOR_MAP.items():
                    node_styles[entity_type] = {
                        "color": color,
                        "borderColor": "#000000",
                        "borderWidth": 1,
                        "caption": "name"
                    }
                
                # 关系样式配置
                rel_styles = {
                    "*": {
                        "color": "#A9A9A9",
                        "width": "weight",
                        "caption": "original_relation"
                    }
                }
                
                # 尝试更新Neo4j Browser的可视化配置
                session.run("""
                CALL apoc.browser.setConfig('browser.visualization.nodeStyleMapping', $nodeStyles)
                """, nodeStyles=node_styles)
                
                session.run("""
                CALL apoc.browser.setConfig('browser.visualization.relationshipStyleMapping', $relStyles)
                """, relStyles=rel_styles)
                
                logging.info("成功设置Neo4j Browser可视化配置")
            except Exception as e:
                logging.warning(f"设置Neo4j Browser可视化配置失败: {str(e)}")
            
            return True
        except Exception as e:
            logging.error(f"导入Neo4j失败: {str(e)}")
            import traceback
            logging.error(traceback.format_exc())
            return False
    
    def _format_relation_name(self, relation: str) -> str:
        """将关系名称格式化为Neo4j兼容的格式"""
        # 保留原始关系名称的大小写，但确保符合Neo4j的命名规则
        formatted = ''.join(c if c.isalnum() or c == '_' else '_' for c in relation)
        
        # 确保不以数字开头
        if formatted and formatted[0].isdigit():
            formatted = 'REL_' + formatted
        
        return formatted

def main():
    """主函数"""
    if len(sys.argv) < 4:
        print("用法: python import_triples_to_neo4j.py <三元组文件路径> <neo4j_uri> <neo4j_user> <neo4j_password>")
        print("示例: python import_triples_to_neo4j.py triples.txt bolt://localhost:7687 neo4j password")
        return
    
    file_path = sys.argv[1]
    neo4j_uri = sys.argv[2]
    neo4j_user = sys.argv[3]
    neo4j_password = sys.argv[4]
    
    if not os.path.exists(file_path):
        print(f"错误: 文件 {file_path} 不存在")
        return
    
    # 开始处理
    start_time = time.time()
    
    try:
        # 创建处理器(必须提供Neo4j连接信息)
        processor = TripleProcessor(
            batch_size=500,
            uri=neo4j_uri,
            user=neo4j_user,
            password=neo4j_password
        )
        
        # 清空数据库并处理三元组
        print("正在处理数据并导入Neo4j...")
        stats = processor.process_triples(file_path)
        
        if stats.get("success", False):
            print(f"处理成功! 用时: {time.time() - start_time:.2f}秒")
            print(f"共处理 {stats['total_triples']} 个三元组")
            print(f"共 {stats['unique_entities']} 个唯一实体，{stats['unique_relationships']} 种关系类型")
            print(f"实体类型统计:")
            for entity_type, count in stats['entity_type_stats'].items():
                print(f"  - {entity_type}: {count}个实体")
            print("数据已成功导入到Neo4j数据库！")
        else:
            print(f"处理失败: {stats.get('error', '未知错误')}")
    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
    finally:
        # 确保关闭数据库连接
        if 'processor' in locals():
            processor.close()

if __name__ == "__main__":
    main() 